package cg.studio.algorithm.tree;

import java.util.Stack;

public class TreeUtil
{
  /**
   * get the most recent common ancestor of two tree nodes
   * @param root
   * @param node1
   * @param node2
   * @return
   */
  public static < T > BinaryTree< T > getCommonAncestorRecursively( BinaryTree< T > root, BinaryTree< T > node1,
                                                                    BinaryTree< T > node2 )
  {
    if( root == null )
      return null;
    if( !contains( root, node1 ) || !contains( root, node2) )
      return null;
    if ( ( node1 == root && contains( root, node2 ) ) || ( node2 == root && contains( root, node1 ) ) )
      return root;
    
    if( contains( root.getLeft(), node1 ) )
    {
      if( contains( root.getRight(), node2 ) )
      {
        //one in left and one in right, this is the ancestor
        return root;
      }
      else
      {
        //both are in left 
        return getCommonAncestorRecursively( root.getLeft(), node1, node2 );
      }
    }
    else    //node1 in right
    {
      if( contains( root.getLeft(), node2 ) )
        return root;
      return getCommonAncestorRecursively( root.getRight(), node1, node2 );
    }
  }
  
  /**
   * check if the tree( root ) contains node(node)
   * @param root
   * @param node
   * @return
   */
  public static < T > boolean contains( BinaryTree< T > root, BinaryTree< T > node )
  {
    if( root == null || node == null )
      return false;
    if( root == node )
      return true;
    
    if( root.getLeft() != null && contains( root.getLeft(), node ) )
      return true;
    if( root.getRight() != null && contains( root.getRight(), node ) )
      return true;
    return false;
  }
  
  /**
   * This method traverse the tree and keep the super nodes in stack, mark the flag when find first node,
   * 
   * @param root
   * @param node1
   * @param node2
   * @return
   */
  public static < T > BinaryTree< T > getCommonAncestorOptimized( BinaryTree< T > root, BinaryTree< T > node1,
                                                                  BinaryTree< T > node2 )
  {
    if( root == null || node1 == null || node2 == null )
      return null;
    
    Stack< BinaryTree<T> > stack = new Stack< BinaryTree<T> >();
    
    boolean matchANode = false;
    BinaryTree<T> anotherNode = null;
    BinaryTree<T> ancestor = null;
    
    while( true )
    {
      //left first
      while( root != null )
      {
        if( !matchANode )
        {
          if( root == node1 || root == node2 )
          {
            matchANode = true;
            anotherNode = ( root != node1 ) ? node1 : node2;
            ancestor = root;
          }
        }
        else
        {
          if( root == anotherNode )
          {
            //both matched
            return ancestor;
          }
        }
        
        if( root.getLeft() != null )
          stack.push( root );
        root = root.getLeft();
      }
      
      //left until leaf, pop to traverse right
      if( stack.isEmpty() )
        return null;
      
      root = stack.pop();
      ancestor = root;
      root = root.getRight();
    }
  }
}
